{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook written by [Zhedong Zheng](https://github.com/zhedongzheng)\n",
    "\n",
    "![title](end2end_mn.png)\n",
    "\n",
    "[End-To-End Memory Networks](https://arxiv.org/abs/1503.08895)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bunch import Bunch\n",
    "from copy import deepcopy\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import json, pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Bunch({\n",
    "    'n_epochs': 20,\n",
    "    'batch_size': 64,\n",
    "    'hidden_dim': 64,\n",
    "    'dropout_rate': 0.3,\n",
    "    'n_hops': 2,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseDataLoader(object):\n",
    "    def __init__(self):\n",
    "        self.data = {\n",
    "            'size': None,\n",
    "            'val':{\n",
    "                'inputs': None,\n",
    "                'questions': None,\n",
    "                'answers': None,},\n",
    "            'len':{\n",
    "                'inputs_len': None,\n",
    "                'inputs_sent_len': None,\n",
    "                'questions_len': None,\n",
    "                'answers_len': None}\n",
    "        }\n",
    "        self.vocab = {\n",
    "            'size': None,\n",
    "            'word2idx': None,\n",
    "            'idx2word': None,\n",
    "        }\n",
    "        self.params = {\n",
    "            'vocab_size': None,\n",
    "            '<start>': None,\n",
    "            '<end>': None,\n",
    "            'max_input_len': None,\n",
    "            'max_sent_len': None,\n",
    "            'max_quest_len': None,\n",
    "            'max_answer_len': None,\n",
    "        }\n",
    "\n",
    "    def input_fn(self):\n",
    "        return tf.estimator.inputs.numpy_input_fn(\n",
    "            x = {\n",
    "                'inputs': self.data['val']['inputs'],\n",
    "                'questions': self.data['val']['questions'],\n",
    "                'inputs_len': self.data['len']['inputs_len'],\n",
    "                'inputs_sent_len': self.data['len']['inputs_sent_len'],\n",
    "                'questions_len': self.data['len']['questions_len'],\n",
    "                'answers_len': self.data['len']['answers_len']\n",
    "            },\n",
    "            y = self.data['val']['answers'] if self.is_training else None,\n",
    "            batch_size = args.batch_size,\n",
    "            num_epochs = args.n_epochs if self.is_training else 1,\n",
    "            shuffle = self.is_training)\n",
    "\n",
    "\n",
    "class DataLoader(BaseDataLoader):\n",
    "    def __init__(self, path, is_training, vocab=None, params=None):\n",
    "        super().__init__()\n",
    "        data, lens = self.load_data(path)\n",
    "        if is_training:\n",
    "            self.build_vocab(data)\n",
    "        else:\n",
    "            self.demo = data\n",
    "            self.vocab = vocab\n",
    "            self.params = deepcopy(params)\n",
    "        self.is_training = is_training\n",
    "        self.padding(data, lens)\n",
    "\n",
    "\n",
    "    def load_data(self, path):\n",
    "        data, lens = bAbI_data_load(path)\n",
    "        self.data['size'] = len(data[0])\n",
    "        return data, lens\n",
    "\n",
    "\n",
    "    def build_vocab(self, data):\n",
    "        signals = ['<pad>', '<unk>', '<start>', '<end>']\n",
    "        inputs, questions, answers = data\n",
    "        i_words = [w for facts in inputs for fact in facts for w in fact if w != '<end>']\n",
    "        q_words = [w for question in questions for w in question]\n",
    "        a_words = [w for answer in answers for w in answer if w != '<end>']\n",
    "        words = list(set(i_words + q_words + a_words))\n",
    "        self.params['vocab_size'] = len(words) + 4\n",
    "        self.params['<start>'] = 2\n",
    "        self.params['<end>'] = 3\n",
    "        self.vocab['word2idx'] = {word: idx for idx, word in enumerate(signals + words)}\n",
    "        self.vocab['idx2word'] = {idx: word for word, idx in self.vocab['word2idx'].items()}\n",
    "        \n",
    "\n",
    "    def padding(self, data, lens):\n",
    "        inputs_len, inputs_sent_len, questions_len, answers_len = lens\n",
    "\n",
    "        self.params['max_input_len'] = max(inputs_len)\n",
    "        self.params['max_sent_len'] = max([fact_len for batch in inputs_sent_len for fact_len in batch])\n",
    "        self.params['max_quest_len'] = max(questions_len)\n",
    "        self.params['max_answer_len'] = max(answers_len)\n",
    "\n",
    "        self.data['len']['inputs_len'] = np.array(inputs_len)\n",
    "        for batch in inputs_sent_len:\n",
    "            batch += [0] * (self.params['max_input_len'] - len(batch))\n",
    "        self.data['len']['inputs_sent_len'] = np.array(inputs_sent_len)\n",
    "        self.data['len']['questions_len'] = np.array(questions_len)\n",
    "        self.data['len']['answers_len'] = np.array(answers_len)\n",
    "        \n",
    "        inputs, questions, answers = deepcopy(data)\n",
    "        for facts in inputs:\n",
    "            for sentence in facts:\n",
    "                for i in range(len(sentence)):\n",
    "                    sentence[i] = self.vocab['word2idx'].get(sentence[i], self.vocab['word2idx']['<unk>'])\n",
    "                sentence += [0] * (self.params['max_sent_len'] - len(sentence))\n",
    "            paddings = [0] * self.params['max_sent_len']\n",
    "            facts += [paddings] * (self.params['max_input_len'] - len(facts))\n",
    "        for question in questions:\n",
    "            for i in range(len(question)):\n",
    "                question[i] = self.vocab['word2idx'].get(question[i], self.vocab['word2idx']['<unk>'])\n",
    "            question += [0] * (self.params['max_quest_len'] - len(question))\n",
    "        for answer in answers:\n",
    "            for i in range(len(answer)):\n",
    "                answer[i] = self.vocab['word2idx'].get(answer[i], self.vocab['word2idx']['<unk>'])\n",
    "\n",
    "        self.data['val']['inputs'] = np.array(inputs)\n",
    "        self.data['val']['questions'] = np.array(questions)\n",
    "        self.data['val']['answers'] = np.array(answers)\n",
    "\n",
    "\n",
    "def bAbI_data_load(path, END=['<end>']):\n",
    "    inputs = []\n",
    "    questions = []\n",
    "    answers = []\n",
    "\n",
    "    inputs_len = []\n",
    "    inputs_sent_len = []\n",
    "    questions_len = []\n",
    "    answers_len = []\n",
    "\n",
    "    for d in open(path):\n",
    "        index = d.split(' ')[0]\n",
    "        if index == '1':\n",
    "            fact = []\n",
    "        if '?' in d:\n",
    "            temp = d.split('\\t')\n",
    "            q = temp[0].strip().replace('?', '').split(' ')[1:] + ['?']\n",
    "            a = temp[1].split() + END\n",
    "            fact_copied = deepcopy(fact)\n",
    "            inputs.append(fact_copied)\n",
    "            questions.append(q)\n",
    "            answers.append(a)\n",
    "\n",
    "            inputs_len.append(len(fact_copied))\n",
    "            inputs_sent_len.append([len(s) for s in fact_copied])\n",
    "            questions_len.append(len(q))\n",
    "            answers_len.append(len(a))\n",
    "        else:\n",
    "            tokens = d.replace('.', '').replace('\\n', '').split(' ')[1:] + END\n",
    "            fact.append(tokens)\n",
    "    return [inputs, questions, answers], [inputs_len, inputs_sent_len, questions_len, answers_len]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode, params):\n",
    "    logits_or_ids = forward(features, labels, mode, params)\n",
    "\n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        return tf.estimator.EstimatorSpec(mode=mode, predictions=logits_or_ids)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        loss_op = tf.reduce_mean(tf.contrib.seq2seq.sequence_loss(\n",
    "            logits = logits_or_ids,\n",
    "            targets = labels,\n",
    "            weights = tf.ones_like(labels, tf.float32)))\n",
    "\n",
    "        train_op = tf.train.AdamOptimizer().minimize(\n",
    "            loss_op,\n",
    "            global_step = tf.train.get_global_step())\n",
    "        return tf.estimator.EstimatorSpec(mode = mode,\n",
    "                                          loss = loss_op,\n",
    "                                          train_op = train_op)\n",
    "\n",
    "\n",
    "def hop_forward(features, question, memory_o, memory_i, response_proj, params, is_training):\n",
    "    match = tf.matmul(question, memory_i, transpose_b=True)\n",
    "\n",
    "    match = pre_softmax_masking(match, features['inputs_len'])\n",
    "\n",
    "    match = tf.nn.softmax(match)       # (batch, question_maxlen, input_maxlen)\n",
    "\n",
    "    match = post_softmax_masking(match, features['questions_len'])\n",
    "\n",
    "    response = tf.matmul(match, memory_o)\n",
    "\n",
    "    response = response_proj(tf.concat([response, question], -1))\n",
    "    \n",
    "    return response\n",
    "\n",
    "\n",
    "def forward(features, labels, mode, params):\n",
    "    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
    "    \n",
    "    with tf.variable_scope('questions'):\n",
    "        question = quest_mem(features['questions'], params, is_training)\n",
    "        \n",
    "    with tf.variable_scope('memory_o'):\n",
    "        memory_o = input_mem(features['inputs'], params, is_training)\n",
    "    \n",
    "    with tf.variable_scope('memory_i'):\n",
    "        memory_i = input_mem(features['inputs'], params, is_training)\n",
    "    \n",
    "    with tf.variable_scope('interaction'):\n",
    "        response_proj = tf.layers.Dense(args.hidden_dim)\n",
    "        \n",
    "        for _ in range(args['n_hops']):\n",
    "            answer = hop_forward(features,\n",
    "                                 question,\n",
    "                                 memory_o,\n",
    "                                 memory_i,\n",
    "                                 response_proj,\n",
    "                                 params,\n",
    "                                 is_training)\n",
    "            question = answer\n",
    "    \n",
    "    with tf.variable_scope('memory_o', reuse=True):\n",
    "        embedding = tf.get_variable('lookup_table')\n",
    "    \n",
    "    with tf.variable_scope('answer'):\n",
    "        output = answer_module(features, params, answer, embedding, is_training, labels)\n",
    "    \n",
    "    return output\n",
    "\n",
    "\n",
    "def input_mem(x, params, is_training):\n",
    "    x = embed_seq(x, params)\n",
    "    x = tf.layers.dropout(x, args.dropout_rate, training=is_training)\n",
    "    pos = position_encoding(params['max_sent_len'], args.hidden_dim)\n",
    "    x = tf.reduce_sum(x * pos, 2)\n",
    "    return x\n",
    "\n",
    "\n",
    "def quest_mem(x, params, is_training):\n",
    "    x = embed_seq(x, params)\n",
    "    x = tf.layers.dropout(x, args.dropout_rate, training=is_training)\n",
    "    pos = position_encoding(params['max_quest_len'], args.hidden_dim)\n",
    "    return (x * pos)\n",
    "\n",
    "\n",
    "def answer_module(features, params, answer, embedding, is_training, labels):\n",
    "    cell = tf.nn.rnn_cell.BasicRNNCell(args.hidden_dim)\n",
    "    vocab_proj = tf.layers.Dense(params['vocab_size'])\n",
    "    state_proj = tf.layers.Dense(args.hidden_dim)\n",
    "    \n",
    "    init_state = state_proj(tf.layers.flatten(answer))\n",
    "    init_state = tf.layers.dropout(init_state, args.dropout_rate, training=is_training)\n",
    "\n",
    "    if is_training:\n",
    "        helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "            inputs = tf.nn.embedding_lookup(embedding, shift_right(labels, params)),\n",
    "            sequence_length = tf.to_int32(features['answers_len']))\n",
    "        decoder = tf.contrib.seq2seq.BasicDecoder(\n",
    "            cell = cell,\n",
    "            helper = helper,\n",
    "            initial_state = init_state,\n",
    "            output_layer = vocab_proj)\n",
    "        decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "            decoder = decoder)\n",
    "        return decoder_output.rnn_output\n",
    "    else:\n",
    "        helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(\n",
    "            embedding = embedding,\n",
    "            start_tokens = tf.tile(\n",
    "                tf.constant([params['<start>']], dtype=tf.int32), [tf.shape(init_state)[0]]),\n",
    "            end_token = params['<end>'])\n",
    "        decoder = tf.contrib.seq2seq.BasicDecoder(\n",
    "            cell = cell,\n",
    "            helper = helper,\n",
    "            initial_state = init_state,\n",
    "            output_layer = vocab_proj)\n",
    "        decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "            decoder = decoder,\n",
    "            maximum_iterations = params['max_answer_len'])\n",
    "        return decoder_output.sample_id\n",
    "\n",
    "\n",
    "def pre_softmax_masking(x, seq_len):\n",
    "    paddings = tf.fill(tf.shape(x), float('-inf'))\n",
    "    T = x.get_shape()[1].value\n",
    "    max_seq_len = x.get_shape()[2].value\n",
    "    masks = tf.sequence_mask(seq_len, max_seq_len, dtype=tf.float32)\n",
    "    masks = tf.tile(tf.expand_dims(masks, 1), [1, T, 1])\n",
    "    return tf.where(tf.equal(masks, 0), paddings, x)\n",
    "\n",
    "\n",
    "def post_softmax_masking(x, seq_len):\n",
    "    T = x.get_shape()[2].value\n",
    "    max_seq_len = x.get_shape()[1].value\n",
    "    masks = tf.sequence_mask(seq_len, max_seq_len, dtype=tf.float32)\n",
    "    masks = tf.tile(tf.expand_dims(masks, -1), [1, 1, T])\n",
    "    return (x * masks)\n",
    "\n",
    "\n",
    "def shift_right(x, params):\n",
    "    batch_size = tf.shape(x)[0]\n",
    "    start = tf.to_int64(tf.fill([batch_size, 1], params['<start>']))\n",
    "    return tf.concat([start, x[:, :-1]], 1)\n",
    "\n",
    "\n",
    "def embed_seq(x, params, zero_pad=True):\n",
    "    lookup_table = tf.get_variable('lookup_table', [params['vocab_size'], args.hidden_dim], tf.float32)\n",
    "    if zero_pad:\n",
    "        lookup_table = tf.concat((tf.zeros([1, args.hidden_dim]), lookup_table[1:, :]), axis=0)\n",
    "    return tf.nn.embedding_lookup(lookup_table, x)\n",
    "\n",
    "\n",
    "def position_encoding(sentence_size, embedding_size):\n",
    "    encoding = np.ones((embedding_size, sentence_size), dtype=np.float32)\n",
    "    ls = sentence_size + 1\n",
    "    le = embedding_size + 1\n",
    "    for i in range(1, le):\n",
    "        for j in range(1, ls):\n",
    "            encoding[i-1, j-1] = (i - (le-1)/2) * (j - (ls-1)/2)\n",
    "    encoding = 1 + 4 * encoding / embedding_size / sentence_size\n",
    "    return np.transpose(encoding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"n_epochs\": 20,\n",
      "    \"batch_size\": 64,\n",
      "    \"hidden_dim\": 64,\n",
      "    \"dropout_rate\": 0.3,\n",
      "    \"n_hops\": 2\n",
      "}\n",
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp1jcnehia\n",
      "INFO:tensorflow:Using config: {'_model_dir': '/var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp1jcnehia', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x121c5c400>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 1 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp1jcnehia/model.ckpt.\n",
      "INFO:tensorflow:loss = 3.9246573, step = 1\n",
      "INFO:tensorflow:global_step/sec: 14.6273\n",
      "INFO:tensorflow:loss = 0.6716565, step = 101 (6.838 sec)\n",
      "INFO:tensorflow:global_step/sec: 16.304\n",
      "INFO:tensorflow:loss = 0.48168907, step = 201 (6.135 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.1168\n",
      "INFO:tensorflow:loss = 0.27430642, step = 301 (6.615 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.7463\n",
      "INFO:tensorflow:loss = 0.2662364, step = 401 (6.350 sec)\n",
      "INFO:tensorflow:global_step/sec: 14.7179\n",
      "INFO:tensorflow:loss = 0.2081354, step = 501 (6.794 sec)\n",
      "INFO:tensorflow:global_step/sec: 14.7563\n",
      "INFO:tensorflow:loss = 0.26827732, step = 601 (6.777 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.5607\n",
      "INFO:tensorflow:loss = 0.22631359, step = 701 (6.426 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.668\n",
      "INFO:tensorflow:loss = 0.235158, step = 801 (6.382 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.7908\n",
      "INFO:tensorflow:loss = 0.21789706, step = 901 (6.333 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.6376\n",
      "INFO:tensorflow:loss = 0.188351, step = 1001 (6.395 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.177\n",
      "INFO:tensorflow:loss = 0.20937227, step = 1101 (6.589 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.9061\n",
      "INFO:tensorflow:loss = 0.13590635, step = 1201 (6.287 sec)\n",
      "INFO:tensorflow:global_step/sec: 13.1814\n",
      "INFO:tensorflow:loss = 0.12292345, step = 1301 (7.586 sec)\n",
      "INFO:tensorflow:global_step/sec: 14.9522\n",
      "INFO:tensorflow:loss = 0.14950338, step = 1401 (6.688 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.6242\n",
      "INFO:tensorflow:loss = 0.22479314, step = 1501 (6.400 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.7117\n",
      "INFO:tensorflow:loss = 0.29873073, step = 1601 (6.365 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.4844\n",
      "INFO:tensorflow:loss = 0.22319978, step = 1701 (6.458 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.1892\n",
      "INFO:tensorflow:loss = 0.1576772, step = 1801 (6.584 sec)\n",
      "INFO:tensorflow:global_step/sec: 14.7724\n",
      "INFO:tensorflow:loss = 0.15976223, step = 1901 (6.769 sec)\n",
      "INFO:tensorflow:global_step/sec: 14.5935\n",
      "INFO:tensorflow:loss = 0.1715068, step = 2001 (6.852 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.6632\n",
      "INFO:tensorflow:loss = 0.1879234, step = 2101 (6.384 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.1615\n",
      "INFO:tensorflow:loss = 0.14548796, step = 2201 (6.595 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.1362\n",
      "INFO:tensorflow:loss = 0.11843179, step = 2301 (6.607 sec)\n",
      "INFO:tensorflow:global_step/sec: 14.3584\n",
      "INFO:tensorflow:loss = 0.15873466, step = 2401 (6.965 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.4072\n",
      "INFO:tensorflow:loss = 0.18857843, step = 2501 (6.490 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.5922\n",
      "INFO:tensorflow:loss = 0.14798367, step = 2601 (6.414 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.9325\n",
      "INFO:tensorflow:loss = 0.20341347, step = 2701 (6.277 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.9719\n",
      "INFO:tensorflow:loss = 0.09403967, step = 2801 (6.261 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.2516\n",
      "INFO:tensorflow:loss = 0.18142506, step = 2901 (6.557 sec)\n",
      "INFO:tensorflow:global_step/sec: 14.6025\n",
      "INFO:tensorflow:loss = 0.11592973, step = 3001 (6.848 sec)\n",
      "INFO:tensorflow:global_step/sec: 15.1873\n",
      "INFO:tensorflow:loss = 0.20867637, step = 3101 (6.584 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 3125 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp1jcnehia/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.14034261.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmp1jcnehia/model.ckpt-3125\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "Testing Accuracy: 0.909\n",
      "\n",
      "[['Fred', 'picked', 'up', 'the', 'football', 'there', '<end>'],\n",
      " ['Fred', 'gave', 'the', 'football', 'to', 'Jeff', '<end>'],\n",
      " ['Bill', 'went', 'back', 'to', 'the', 'bathroom', '<end>'],\n",
      " ['Jeff', 'grabbed', 'the', 'milk', 'there', '<end>'],\n",
      " ['Jeff', 'gave', 'the', 'football', 'to', 'Fred', '<end>'],\n",
      " ['Fred', 'handed', 'the', 'football', 'to', 'Jeff', '<end>'],\n",
      " ['Jeff', 'handed', 'the', 'football', 'to', 'Fred', '<end>'],\n",
      " ['Fred', 'gave', 'the', 'football', 'to', 'Jeff', '<end>']]\n",
      "\n",
      "Question: ['Who', 'did', 'Fred', 'give', 'the', 'football', 'to', '?']\n",
      "\n",
      "Prediction: ['Jeff', '<end>']\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    tf.logging.set_verbosity(tf.logging.INFO)\n",
    "    print(json.dumps(args, indent=4))\n",
    "\n",
    "    train_dl = DataLoader(\n",
    "        path='../temp/qa5_three-arg-relations_train.txt',\n",
    "        is_training=True)\n",
    "    test_dl = DataLoader(\n",
    "        path='../temp/qa5_three-arg-relations_test.txt',\n",
    "        is_training=False, vocab=train_dl.vocab, params=train_dl.params)\n",
    "\n",
    "    model = tf.estimator.Estimator(model_fn, params=train_dl.params)\n",
    "    model.train(train_dl.input_fn())\n",
    "    gen = model.predict(test_dl.input_fn())\n",
    "    preds = np.concatenate(list(gen))\n",
    "    preds = np.reshape(preds, [test_dl.data['size'], 2])\n",
    "    print('Testing Accuracy:', (test_dl.data['val']['answers'][:, 0] == preds[:, 0]).mean())\n",
    "    demo(test_dl.demo, test_dl.vocab['idx2word'], preds)\n",
    "\n",
    "\n",
    "def demo(demo, idx2word, ids, demo_idx=3):\n",
    "    demo_i, demo_q, demo_a = demo\n",
    "    print()\n",
    "    pprint.pprint(demo_i[demo_idx])\n",
    "    print()\n",
    "    print('Question:', demo_q[demo_idx])\n",
    "    print()\n",
    "    print('Prediction:', [idx2word[id] for id in ids[demo_idx]])\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
